#include <core.p4>
#include <v1model.p4>

const bit<16> ENCAP_PROTO_ID = 0x1212;
const bit<16> ETH_TYPE_IP = 0x800;

typedef bit<9>  PortId_t;
typedef bit<48> MacAddress_t;
typedef bit<32> IpAddress_t;

header L2_Header {
    MacAddress_t dest_mac;
    MacAddress_t src_mac;
    bit<16>    next_proto;
}

header CustomOverlay_Hdr {
    bit<16> inner_type;
    bit<16> tunnel_id;
}

header L3_IPv4_Header {
    bit<4>    ver;
    bit<4>    hdr_len;
    bit<8>    dscp;
    bit<16>    total_len;
    bit<16>    identification;
    bit<3>    flags;
    bit<13>    frag_offset;
    bit<8>    time_to_live;
    bit<8>    protocol;
    bit<16>    chksum;
    IpAddress_t source_ip;
    IpAddress_t destination_ip;
}

struct internal_metadata {
    bit<1> is_tunnel_pkt;
    bit<1> needs_encap;
    bit<1> needs_decap;
    PortId_t out_port_temp;
}

struct PacketHeaders {
    L2_Header    eth;
    CustomOverlay_Hdr    tunnel;
    L3_IPv4_Header    ip;
}

parser FlowParser(packet_in pkt_in,
                out PacketHeaders pkt_hdr,
                inout internal_metadata local_meta,
                inout standard_metadata_t std_meta) {

    state start {
        transition extract_l2;
    }

    state extract_l2 {
        pkt_in.extract(pkt_hdr.eth);
        transition select(pkt_hdr.eth.next_proto) {
            ENCAP_PROTO_ID: extract_tunnel;
            ETH_TYPE_IP: extract_ipv4;
            default: accept;
        }
    }

    state extract_tunnel {
        pkt_in.extract(pkt_hdr.tunnel);
        transition select(pkt_hdr.tunnel.inner_type) {
            ETH_TYPE_IP: extract_ipv4;
            default: accept;
        }
    }

    state extract_ipv4 {
        pkt_in.extract(pkt_hdr.ip);
        transition accept;
    }
}

control Verify_Checksum(inout PacketHeaders pkt_hdr, inout internal_metadata local_meta) {
    apply {  }
}

counter(1024, CounterType.packets) Encapsulation_Counter;
counter(1024, CounterType.packets) Decapsulation_Counter;

control IngressProcessing(inout PacketHeaders pkt_hdr,
                  inout internal_metadata local_meta,
                  inout standard_metadata_t std_meta) {

    action DropPacket() {
        mark_to_drop(std_meta);
    }

    action Process_L3_Forwarding(MacAddress_t next_hop_mac, PortId_t out_port) {
        std_meta.egress_spec = out_port;
        pkt_hdr.eth.src_mac = pkt_hdr.eth.dest_mac;
        pkt_hdr.eth.dest_mac = next_hop_mac;
        pkt_hdr.ip.time_to_live = pkt_hdr.ip.time_to_live - 1;
        local_meta.needs_encap = 0;
    }

    action Prepare_Encapsulation(bit<16> dest_tunnel_id) {
        local_meta.needs_encap = 1;
        local_meta.is_tunnel_pkt = 1;
        pkt_hdr.tunnel.tunnel_id = dest_tunnel_id;
    }

    action Setup_Decapsulation(MacAddress_t final_dest_mac, PortId_t out_port) {
        local_meta.needs_decap = 1;
        local_meta.out_port_temp = out_port;
        pkt_hdr.eth.dest_mac = final_dest_mac;
    }
    
    action Tunnel_Pass_Through(PortId_t out_port) {
        std_meta.egress_spec = out_port;
        local_meta.needs_decap = 0;
    }

    table IP_Destination_Lookup {
        key = {
            pkt_hdr.ip.destination_ip: lpm;
        }
        actions = {
            Process_L3_Forwarding;
            Prepare_Encapsulation;
            DropPacket;
            NoAction;
        }
        size = 1024;
        default_action = DropPacket();
    }

    table Overlay_ID_Match {
        key = {
            pkt_hdr.tunnel.tunnel_id: exact;
        }
        actions = {
            Tunnel_Pass_Through;
            Setup_Decapsulation;
            DropPacket;
        }
        size = 1024;
        default_action = DropPacket();
    }

    apply {
        local_meta.is_tunnel_pkt = pkt_hdr.tunnel.isValid();
        local_meta.needs_encap = 0;
        local_meta.needs_decap = 0;
        
        if (local_meta.is_tunnel_pkt == 1) {
            Overlay_ID_Match.apply();

            if (local_meta.needs_decap == 1) {
                std_meta.egress_spec = local_meta.out_port_temp;
                pkt_hdr.eth.next_proto = pkt_hdr.tunnel.inner_type;
                pkt_hdr.tunnel.setInvalid();
                Decapsulation_Counter.count((bit<32>)pkt_hdr.tunnel.tunnel_id);
            }
        } 
        else if (pkt_hdr.ip.isValid() == 1) {
            IP_Destination_Lookup.apply();
        }
        
        if (local_meta.needs_encap == 1) {
            pkt_hdr.tunnel.setValid();
            pkt_hdr.tunnel.inner_type = pkt_hdr.eth.next_proto;
            pkt_hdr.eth.next_proto = ENCAP_PROTO_ID;
            Encapsulation_Counter.count((bit<32>)pkt_hdr.tunnel.tunnel_id);
            
            Overlay_ID_Match.apply();
            local_meta.needs_decap = 0;
        }
    }
}

control EgressProcessing(inout PacketHeaders pkt_hdr,
                 inout internal_metadata local_meta,
                 inout standard_metadata_t std_meta) {
    apply {  }
}

control Calculate_Checksum(inout PacketHeaders pkt_hdr, inout internal_metadata local_meta) {
     apply {
        update_checksum(
            pkt_hdr.ip.isValid(),
            { pkt_hdr.ip.ver,
              pkt_hdr.ip.hdr_len,
              pkt_hdr.ip.dscp,
              pkt_hdr.ip.total_len,
              pkt_hdr.ip.identification,
              pkt_hdr.ip.flags,
              pkt_hdr.ip.frag_offset,
              pkt_hdr.ip.time_to_live,
              pkt_hdr.ip.protocol,
              pkt_hdr.ip.source_ip,
              pkt_hdr.ip.destination_ip },
            pkt_hdr.ip.chksum,
            HashAlgorithm.csum16);
    }
}

control DataSerializer(packet_out pkt_out, in PacketHeaders pkt_hdr) {
    apply {
        pkt_out.emit(pkt_hdr.eth);
        pkt_out.emit(pkt_hdr.tunnel);
        pkt_out.emit(pkt_hdr.ip);
    }
}

V1Switch(
FlowParser(),
Verify_Checksum(),
IngressProcessing(),
EgressProcessing(),
Calculate_Checksum(),
DataSerializer()
) core_switch_main;